import numpy
from matplotlib import pyplot as plt
from PIL import Image

numpy.seterr(divide='ignore', invalid='ignore')


def my_SIFT_process(original_img, harris_res_list):
    # sobel_x = numpy.array([[-3, 0, 3], [-10, 0, 10], [-3, 0, 3]],
    #                       dtype=numpy.float32)
    # sobel_y = numpy.array([[-3, -10, -3], [0, 0, 0], [3, 10, 3]],
    #                       dtype=numpy.float32)
    # laplas_x = numpy.array([[1, -2, 1], [2, -4, 2], [1, -2, 1]],
    #                        dtype=numpy.float32)
    scale_cube = LoG_creator(original_img)
    after_scale_set = set()
    # 这里已经求出两个scale下x方向和y方向的梯度
    d_f1_x, d_f1_y = sobel_f(scale_cube[1])
    d_f2_x, d_f2_y = sobel_f(scale_cube[2])
    # d_f2_x = my_convolution(scale_cube[2], sobel_x)
    # d_f2_y = my_convolution(scale_cube[2], sobel_y)
    d_func = {1: (d_f1_x ** 2, d_f1_y ** 2, d_f1_x * d_f1_y), 2: (d_f2_x ** 2, d_f2_y ** 2, d_f2_x * d_f2_y)}
    plt.gray()
    plt.imshow(original_img)
    k = 0
    D_res = numpy.zeros((d_func[1][0].shape[0], d_func[1][0].shape[1]), dtype=numpy.float)
    harris_res_list1 = set()

    do_DoG1, help1 = sobel_f(d_f1_x)
    lo_DoG1 = laplas_f(d_f1_x)
    do_DoG2, help2 = sobel_f(d_f2_x)
    lo_DoG2 = laplas_f(d_f2_x)

    func_D_factor = numpy.array([do_DoG1, lo_DoG1, do_DoG2, lo_DoG2], dtype=numpy.float32)

    for i in range(1, 2):
        for dot in harris_res_list:
            if dot[0] - 1 >= 0 and dot[0] + 2 < original_img.shape[0] and dot[
                1] - 1 >= 0 and dot[1] + 2 < original_img.shape[1]:
                if (scale_cube[i][dot[0], dot[1]] >= (
                        scale_cube[i - 1:i + 2][dot[0] - 1:dot[0] + 2, dot[1] - 1:dot[1] + 2]).all()) or (
                        scale_cube[i][dot[0], dot[1]] <= (
                        scale_cube[i - 1:i + 2][dot[0] - 1:dot[0] + 2, dot[1] - 1:dot[1] + 2]).all()):
                    d_1 = (scale_cube[i][dot[0] + 1, dot[1]] - scale_cube[i][dot[0] - 1, dot[1]]) / (
                            scale_cube[i][dot[0], dot[1] + 1] - scale_cube[i][dot[0], dot[1] - 1])
                    d_2 = (scale_cube[i][dot[0] + 1, dot[1]] + scale_cube[i][dot[0] - 1, dot[1]] - 2 * scale_cube[i][
                        dot[0], dot[1]]) / (scale_cube[i][dot[0], dot[1] + 1] + scale_cube[i][dot[0], dot[1] - 1] - 2 *
                                            scale_cube[i][dot[0], dot[1]])
                    t = -d_1 / d_2
                    val = scale_cube[i][dot[0], dot[1]] + t * func_D_factor[2 * i - 2][dot[0], dot[1]] + (t ** 2) * \
                          func_D_factor[2 * i - 1][dot[0], dot[1]]
                    if abs(val) > 0.03:
                        harris_res_list1.add((dot[0], dot[1], i))

    for i in range(1, 2):
        for dot in harris_res_list1:
            # if scale_cube[i, dot[0], dot[1]] >= numpy.max(scale_cube[i - 1:i + 2, dot[0] - 1:dot[0] + 2,
            # dot[1] - 1:dot[1] + 2]) or scale_cube[i, dot[0], dot[1]] <= numpy.min(scale_cube[i - 1:i + 2,
            # dot[0] - 1:dot[0] + 2, dot[1] - 1:dot[1] + 2]):
            TrH = d_func[i][0][dot[0], dot[1]] + d_func[i][1][dot[0], dot[1]]
            DetH = d_func[i][0][dot[0], dot[1]] * d_func[i][1][dot[0], dot[1]] - (d_func[i][2][dot[0], dot[1]] ** 2)
            if (TrH ** 2 / DetH) <= 12.1:
                if ((dot[0], dot[1], 1) not in after_scale_set) and ((dot[0], dot[1], 2) not in after_scale_set):
                    k += 1
                    after_scale_set.add((dot[0], dot[1], i))
                    plt.scatter([dot[1]], [dot[0]], color='b')
            else:
                D_res[dot[0], dot[1]] = 255
    print(k)
    print(after_scale_set.__len__())

    plt.title('After_edge_cleaning image')  # 标题
    plt.show()

    # 下面要算出每个点的主要方向，以及对应的特征向量。特征向量是选16*16然后4*4作为一组求组内形成的一个8位的向量，代表八个方向，在相加时要先和特征点的主要方向相减，这样算出的角度是相对角度。
    # 计算主要方向需要计算L(x,y)，这个是根据不同sigma进行高斯模糊得到，a层的dog对应的sigma就是两个模糊后相减的图像中减图的sigma
    # 根据这个计算出两层每个点的magnitude x&y 和arctan
    # 然后首先根据特征点周围4*4分成八个区域求和找到主要方向
    # 然后对周围16*16，每个4*4角度相对主要方向先减一下，然后分成八个区域相加，数值形成八位向量，最后得到二维向量：8个一维，每个一维有16位

    layer1 = numpy.zeros((scale_cube[1].shape[0], scale_cube[1].shape[1]), dtype=numpy.float)
    layer2 = numpy.zeros((scale_cube[1].shape[0], scale_cube[1].shape[1]), dtype=numpy.float)
    layer3 = numpy.zeros((scale_cube[1].shape[0], scale_cube[1].shape[1]), dtype=numpy.float)
    layer4 = numpy.zeros((scale_cube[1].shape[0], scale_cube[1].shape[1]), dtype=numpy.float)

    data_layers = [layer1, layer2, layer3, layer4]
    # 两层经过高斯模糊的图片，对应着点的两种不同scale
    # gaussian_blur = [my_convolution(original_img, gaussian_kernel(5, 3 * 2 * 1.5)),
    #                  my_convolution(original_img, gaussian_kernel(5, 3 * numpy.sqrt(2) * 1.5))]
    gaussian_blur = [gaussian_filter(original_img, 5, 1.6 * 2 * 1.5),
                     gaussian_filter(original_img, 5, 1.6 * numpy.sqrt(2) * 1.5)]

    # 为两层每一层都算好arctan和梯度，为下面做准备。
    for x in range(1, gaussian_blur[1].shape[1] - 1):
        for y in range(1, gaussian_blur[1].shape[0] - 1):
            # 被迫分离，0，2是梯度，1，3是角度
            delta_x1 = gaussian_blur[0][y, x + 1] - gaussian_blur[0][y, x - 1]
            delta_y1 = gaussian_blur[0][y + 1, x] - gaussian_blur[0][y - 1, x]
            layer1[y, x] = numpy.sqrt(delta_y1 ** 2 + delta_x1 ** 2)
            layer2[y, x] = numpy.arctan(numpy.divide(delta_y1, delta_x1))
            delta_x2 = gaussian_blur[1][y, x + 1] - gaussian_blur[1][y, x - 1]
            delta_y2 = gaussian_blur[1][y + 1, x] - gaussian_blur[1][y - 1, x]
            layer3[y, x] = numpy.sqrt(delta_y2 ** 2 + delta_x2 ** 2)
            layer4[y, x] = numpy.arctan(numpy.divide(delta_y2, delta_x2))

    # 下面计算出一个点的主要方向
    p_with_orientation = set()
    for interest_point in after_scale_set:
        if interest_point[0] - 2 >= 0 and interest_point[0] + 2 < original_img.shape[0] and interest_point[
            1] - 2 >= 0 and interest_point[1] + 2 < original_img.shape[1]:
            vector_list = numpy.zeros(36)
            matrix_sita = gaussian_filter(
                data_layers[interest_point[2] * 2 - 2][interest_point[0] - 2:interest_point[0] + 2,
                interest_point[1] - 2:interest_point[1] + 2], 3, 3)
            for y_scale in range(interest_point[0] - 2, interest_point[0] + 2):
                for x_scale in range(interest_point[1] - 2, interest_point[1] + 2):
                    how = numpy.floor((data_layers[interest_point[2] * 2 - 1][y_scale, x_scale] / numpy.pi * 18))
                    int_changer = how.astype(int) % 36
                    vector_list[int_changer] += matrix_sita[
                        y_scale - interest_point[0] + 2, x_scale - interest_point[1] + 2]
            print(vector_list)
            main_angel = list_max(vector_list) * (numpy.pi / 18) + numpy.pi / 36
            p_with_orientation.add((interest_point[0], interest_point[1], interest_point[2], main_angel))

    # 16*16 范围下特征向量
    interest_point_set = {}
    k = 0
    for interest_point in p_with_orientation:

        if interest_point[0] - 7 >= 0 and interest_point[0] + 8 < original_img.shape[0] and interest_point[
            1] - 7 >= 0 and interest_point[1] + 8 < original_img.shape[1]:
            res1_mag, res1_arct = gaussian_trans_derivative_filter(
                original_img[interest_point[0] - 7:interest_point[0] + 9, interest_point[1] - 7:interest_point[1] + 9],
                5, 1.6 * 2 * 1.5, interest_point[3])
            res2_mag, res2_arct = gaussian_trans_derivative_filter(
                original_img[interest_point[0] - 7:interest_point[0] + 9, interest_point[1] - 7:interest_point[1] + 9],
                5, 1.6 * numpy.sqrt(2) * 1.5,
                interest_point[3])
            data_layers1 = [res1_mag, res1_arct, res2_mag, res2_arct]
            eigen_vector = []
            for y_scale in range(0, 16, 4):
                for x_scale in range(0, 16, 4):
                    vector_part = [0, 0, 0, 0, 0, 0, 0, 0]
                    for x in range(x_scale, x_scale + 4):
                        for y in range(y_scale, y_scale + 4):
                            # 要先找到对应scale的数据层
                            this_position = numpy.divide(
                                data_layers1[interest_point[2] * 2 - 1][y, x],
                                (numpy.pi / 4))
                            why = numpy.floor(this_position)
                            int_changer = (why.astype(int)) % 8
                            vector_part[int_changer] += data_layers1[interest_point[2] * 2 - 2][y, x]
                    # 将部分向量连接在主向量后
                    for item in vector_part:
                        eigen_vector.append(item)
            # 将特征向量连通特征点加入

            eigen_vector1 = eigen_vector / numpy.linalg.norm(eigen_vector)
            interest_point_set[(interest_point[0], interest_point[1])] = eigen_vector1
            k += 1
    return interest_point_set


def SIFT_match(pic_1, eigen_vectors1, pic2, eigen_vectors2):
    pair_set = list()
    k = 0
    for dots in eigen_vectors1:
        first_match = numpy.inf
        first_dot = dots
        target = ()
        second_match = numpy.inf
        for target_dots in eigen_vectors2:
            dist1 = numpy.linalg.norm(eigen_vectors1[dots] - eigen_vectors2[target_dots])
            if dist1 < first_match:
                second_match = first_match
                first_match = dist1
                first_dot = dots
                target = target_dots
            else:
                if dist1 < second_match:
                    second_match = dist1
        if numpy.divide(first_match, second_match) < 0.8:
            pair_set.append((first_dot[1], first_dot[0], target[1], target[0]))
            k += 1
    print("结果是：")
    print(k)
    print(pair_set)
    return pair_set


def LoG_creator(original_img):
    # img_lv0 = my_convolution(original_img, gaussian_kernel(5, 1.6))
    # img_lv1 = my_convolution(original_img, gaussian_kernel(5, 1.6 * numpy.sqrt(2)))
    # img_lv2 = my_convolution(original_img, gaussian_kernel(5, 1.6 * 2))
    # img_lv3 = my_convolution(original_img, gaussian_kernel(5, 1.6 * 2 * numpy.sqrt(2)))
    # img_lv4 = my_convolution(original_img, gaussian_kernel(5, 1.6 * 4))
    img_lv0 = gaussian_filter(original_img, 5, 1.6)
    img_lv1 = gaussian_filter(original_img, 5, 1.6 * numpy.sqrt(2))
    img_lv2 = gaussian_filter(original_img, 5, 1.6 * 2)
    img_lv3 = gaussian_filter(original_img, 5, 1.6 * 2 * numpy.sqrt(2))
    img_lv4 = gaussian_filter(original_img, 5, 1.6 * 4)
    LoG_0 = img_lv1 - img_lv0
    LoG_1 = img_lv2 - img_lv1
    LoG_2 = img_lv3 - img_lv2
    LoG_3 = img_lv4 - img_lv3
    plt.figure()  # 打开matplotlib的可视化figure

    return numpy.array([LoG_3, LoG_2, LoG_1, LoG_0], dtype=numpy.float32)


def draw1(img1, img2, pair):
    size1, size2 = img1.size, img2.size
    print(size1)
    joint = Image.new("RGB", (max(size1[0], size2[0]), size1[1] + size2[1]))
    loc1, loc2 = (0, 0), (0, size1[1])
    joint.paste(img1, loc1)
    joint.paste(img2, loc2)
    joint.save('Part2Pics/joint.png')
    im = numpy.asarray(Image.open(r'Part2Pics/joint.png'))
    plt.imshow(im)
    for pair_dot in pair:
        plt.plot([pair_dot[0], pair_dot[2]], [pair_dot[1], pair_dot[3] + size1[1]], color='r')
        #  plt.scatter([pair_dot[0], pair_dot[2]], [pair_dot[1], pair_dot[3] + size1[1]], color='b')
    plt.show()


def sobel_f(gray):
    height, weight = gray.shape

    sobel_y = numpy.array(((3, 10, 3), (0, 0, 0), (-3, -10, -3)), numpy.float32)
    sobel_x = numpy.array(((3, 0, -3), (10, 0, -10), (3, 0, -3)), numpy.float32)

    tmp = numpy.pad(gray, (1, 1), 'edge')

    Ix = numpy.zeros_like(gray, numpy.float32)
    Iy = numpy.zeros_like(gray, numpy.float32)

    for y in range(height):
        for x in range(weight):
            Ix[y, x] = numpy.mean(tmp[y: y + 3, x: x + 3] * sobel_x)
            Iy[y, x] = numpy.mean(tmp[y: y + 3, x: x + 3] * sobel_y)

    return Ix, Iy


def laplas_f(gray):
    height, weight = gray.shape

    laplas_x = numpy.array([[1, -2, 1], [2, -4, 2], [1, -2, 1]], numpy.float32)

    tmp = numpy.pad(gray, (1, 1), 'edge')

    Ix = numpy.zeros_like(gray, numpy.float32)

    for y in range(height):
        for x in range(weight):
            Ix[y, x] = numpy.mean(tmp[y: y + 3, x: x + 3] * laplas_x)

    return Ix


def gaussian_filter(I, k=3, sigma=5.0):
    print(I.shape)
    H, W = I.shape
    res = numpy.zeros((H, W), dtype=numpy.float)

    I_t = numpy.pad(I, (k // 2, k // 2), 'edge')

    # gaussian kernel
    K = numpy.zeros((k, k), numpy.float)
    for x in range(k):
        for y in range(k):
            _x = x - k // 2
            _y = y - k // 2
            K[y, x] = numpy.exp(-(_x ** 2 + _y ** 2) / (2 * (sigma ** 2)))
    K /= (sigma * numpy.sqrt(2 * numpy.pi))
    K /= K.sum()

    # filtering
    for y in range(H):
        for x in range(W):
            res[y, x] = numpy.sum(I_t[y: y + k, x: x + k] * K)

    return res


def list_max(list_in):
    max_in = max(list_in)
    a = 0
    for i in range(list_in.size):
        if list_in[i] == max_in:
            a = i
            return a


def trans_to_gray(img):
    gray = 0.2126 * img[..., 2] + 0.7152 * img[..., 1] + 0.0722 * img[..., 0]
    gray = gray.astype(numpy.uint8)
    return gray


def gaussian_trans_derivative_filter(I, k=3, sigma=5.0, xita=0.0):
    H, W = I.shape
    res_mag, res_arctan = numpy.zeros((H, W), dtype=numpy.float), numpy.zeros((H, W), dtype=numpy.float)

    I_t = numpy.pad(I, (k // 2, k // 2), 'edge')

    # gaussian kernel
    K_x, K_y = numpy.zeros((k, k), numpy.float), numpy.zeros((k, k), numpy.float)
    for x in range(k):
        for y in range(k):
            _x0 = x - k // 2
            _y0 = y - k // 2
            _x = _x0 * numpy.cos(xita) - _y0 * numpy.sin(xita)
            _y = _x0 * numpy.sin(xita) + _y0 * numpy.cos(xita)
            K_x[y, x] = numpy.exp(-(_x ** 2 + _y ** 2) / (2 * (sigma ** 2))) * _x * (
                    -_x0 * numpy.sin(xita) - _y0 * numpy.cos(xita)) / -sigma ** 2
            K_y[y, x] = numpy.exp(-(_x ** 2 + _y ** 2) / (2 * (sigma ** 2))) * _y * (
                        _x0 * numpy.cos(xita) - _y0 * numpy.sin(xita)) / -sigma ** 2
    K_x /= (sigma * numpy.sqrt(2 * numpy.pi))
    K_x /= K_x.sum()
    K_y /= (sigma * numpy.sqrt(2 * numpy.pi))
    K_y /= K_y.sum()

    # filtering
    for y in range(H):
        for x in range(W):
            res_mag[y, x] = numpy.sqrt(
                numpy.sum(I_t[y: y + k, x: x + k] * K_x) ** 2 + numpy.sum(I_t[y: y + k, x: x + k] * K_y) ** 2)
            res_arctan[y, x] = numpy.sum(I_t[y: y + k, x: x + k] * K_y) / numpy.sum(I_t[y: y + k, x: x + k] * K_x)

    return res_mag, res_arctan
